package com.omni.monitor.plugin.trace_monitor;

import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.omni.monitor.model.OmniConst;
import com.omni.monitor.model.trace_monitor.ServiceTraceNode;
import com.omni.monitor.model.trace_monitor.TraceNode;
import com.omni.monitor.plugin.utils.SpringContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * 增强版：自动关联父子线程调用链 + 任务后清理
 */
@Slf4j
public class AutoCleanTtlThreadPoolDecorator extends ThreadPoolExecutor {
    private final ThreadPoolExecutor delegate;

    public AutoCleanTtlThreadPoolDecorator(ThreadPoolExecutor delegate) {
        super(
                delegate.getCorePoolSize(),
                delegate.getMaximumPoolSize(),
                delegate.getKeepAliveTime(TimeUnit.NANOSECONDS),
                TimeUnit.NANOSECONDS,
                delegate.getQueue(),
                delegate.getThreadFactory(),
                delegate.getRejectedExecutionHandler()
        );
        this.delegate = delegate;
    }

    // -------------------------- 重写任务提交方法（核心增强） --------------------------
    @Override
    public void execute(Runnable command) {
        delegate.execute(wrapTask(command));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return delegate.submit(wrapTask(task));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        return delegate.submit(wrapTask(task), result);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return delegate.submit(wrapTask(task));
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
        return delegate.invokeAll(wrapTasks(tasks));
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
        return delegate.invokeAll(wrapTasks(tasks), timeout, unit);
    }

    // -------------------------- 核心：增强任务包装逻辑 --------------------------

    /**
     * 包装Runnable：
     * 1. 继承父线程上下文（TTL）
     * 2. 子线程创建"线程切换节点"，关联父线程栈顶
     * 3. 任务执行后自动出栈并清理
     */
    private Runnable wrapTask(Runnable runnable) {
        // 1. 提交任务时立即捕获父节点和traceId（关键修改）
        TraceContext traceContext = SpringContextHolder.getBeanSafe(TraceContext.class);
        TraceNode parentNode = traceContext != null ? traceContext.peekNode() : null;
        String traceId = MDC.get(OmniConst.TRACE_ID);
        // 缓存父节点的快照信息（避免后续变化）
        String parentSpanId = parentNode != null ? parentNode.getSpanId() : "";
        int parentDepth = parentNode != null ? parentNode.getDepth() : 0;
        // 2. TTL包装（继承父线程上下文到子线程）
        Runnable ttlRunnable = TtlRunnable.get(runnable);
        // 在wrapTask方法中，提交任务时缓存MDC信息
        Map<String, String> parentMdc = MDC.getCopyOfContextMap();
        // 3. 二次包装：添加线程切换节点管理
        return () -> {
            TraceNode threadSwitchNode = null;
            try {
                // 恢复父线程MDC上下文
                if (parentMdc != null) {
                    MDC.setContextMap(parentMdc);
                }
                // 4. 校验TTL栈的父节点（若有）
                if (traceContext != null && traceContext.peekNode() != null) {
                    TraceNode ttlParent = traceContext.peekNode();
                    if (!ttlParent.getSpanId().equals(parentSpanId)) {
                        log.warn("TTL栈与缓存父节点不一致：缓存={}，TTL栈={}", parentSpanId, ttlParent.getSpanId());
                        // 强制重置TTL栈为缓存的父节点快照
                        traceContext.forceClear();
                        Deque<TraceNode> newStack = new ConcurrentLinkedDeque<>();
                        if (parentNode != null) {
                            newStack.addFirst(parentNode);
                            traceContext.getTraceStack().addAll(newStack); // 需TraceContext暴露getTraceStack方法
                        }
                    }
                }
                if (traceContext != null && StringUtils.isNotBlank(traceId)) {
                    // a. 创建"线程切换节点"（标记从父线程到子线程的切换）
                    // 使用提交时缓存的父节点信息创建切换节点（关键修改）
                    threadSwitchNode = createThreadSwitchNode(parentSpanId, parentDepth, traceId);
                    // b. 子线程入栈（作为子线程调用链的根节点）
                    traceContext.pushNode(threadSwitchNode);
                    log.debug("[线程切换] 父线程={} → 子线程={}，切换节点spanId={}",
                            parentNode != null ? parentNode.getThreadName() : "null",
                            Thread.currentThread().getName(),
                            threadSwitchNode.getSpanId());
                }
                // c. 执行原始任务（子线程内部的调用链会基于切换节点继续入栈）
                ttlRunnable.run();
            } finally {
                // 清理子线程MDC
                MDC.clear();
                // 4. 任务结束：将桥接节点进行出栈 + 清理子线程上下文
                if (traceContext != null && threadSwitchNode != null) {
                    TraceNode popped = traceContext.popNode(threadSwitchNode.getSpanId());
                    if (popped != null && !popped.getSpanId().equals(threadSwitchNode.getSpanId())) {
                        log.error("[线程切换异常] 出栈节点不匹配，预期={}，实际={}",
                                threadSwitchNode.getSpanId(), popped.getSpanId());
                    }
                    if (traceContext.peekNode() == null) {
                        traceContext.forceClear();
                        log.debug("线程[{}]任务执行完毕，自动清理上下文", Thread.currentThread().getName());
                    }
                }
            }
        };
    }

    /**
     * 包装Callable（逻辑与Runnable一致）
     */
    private <T> Callable<T> wrapTask(Callable<T> callable) {
        TraceContext traceContext = SpringContextHolder.getBeanSafe(TraceContext.class);
        TraceNode parentNode = traceContext != null ? traceContext.peekNode() : null;
        String traceId = MDC.get(OmniConst.TRACE_ID);
        // 缓存父节点的快照信息（避免后续变化）
        String parentSpanId = parentNode != null ? parentNode.getSpanId() : "";
        int parentDepth = parentNode != null ? parentNode.getDepth() : 0;
        Callable<T> ttlCallable = TtlCallable.get(callable);
        // 在wrapTask方法中，提交任务时缓存MDC信息
        Map<String, String> parentMdc = MDC.getCopyOfContextMap();
        return () -> {
            TraceNode threadSwitchNode = null;
            try {
                // 恢复父线程MDC上下文
                if (parentMdc != null) {
                    MDC.setContextMap(parentMdc);
                }
                // 4. 校验TTL栈的父节点（若有）
                if (traceContext != null && traceContext.peekNode() != null) {
                    TraceNode ttlParent = traceContext.peekNode();
                    if (!ttlParent.getSpanId().equals(parentSpanId)) {
                        log.warn("TTL栈与缓存父节点不一致：缓存={}，TTL栈={}", parentSpanId, ttlParent.getSpanId());
                        // 强制重置TTL栈为缓存的父节点快照
                        traceContext.forceClear();
                        Deque<TraceNode> newStack = new ConcurrentLinkedDeque<>();
                        if (parentNode != null) {
                            newStack.addFirst(parentNode);
                            traceContext.getTraceStack().addAll(newStack); // 需TraceContext暴露getTraceStack方法
                        }
                    }
                }

                // 5. 创建并入栈桥接节点
                if (traceContext != null && StringUtils.isNotBlank(traceId)) {
                    threadSwitchNode = createThreadSwitchNode(parentSpanId, parentDepth, traceId);
                    traceContext.pushNode(threadSwitchNode);
                }
                return ttlCallable.call();
            } finally {
                // 清理子线程MDC
                MDC.clear();
                // 4. 任务结束：出栈切换节点 + 清理子线程上下文
                if (traceContext != null && threadSwitchNode != null) {
                    TraceNode popped = traceContext.popNode(threadSwitchNode.getSpanId());
                    if (popped != null && !popped.getSpanId().equals(threadSwitchNode.getSpanId())) {
                        log.error("[线程切换异常] 出栈节点不匹配，预期={}，实际={}",
                                threadSwitchNode.getSpanId(), popped.getSpanId());
                    }
                    if (traceContext.peekNode() == null) {
                        traceContext.forceClear();
                        log.debug("线程[{}]任务执行完毕，自动清理上下文", Thread.currentThread().getName());
                    }
                }
            }
        };
    }

    // 新增辅助方法：基于缓存的父节点信息创建切换节点
    private TraceNode createThreadSwitchNode(String parentSpanId, int parentDepth, String traceId) {
        ServiceTraceNode switchNode = new ServiceTraceNode();
        switchNode.setTraceId(traceId);
        switchNode.setParentSpanId(parentSpanId);
        //桥接节点的spanId与父节点一致，在构建树结构时，会自动将桥接节点的spanId与父节点的spanId关联起来
        switchNode.setSpanId(parentSpanId);
        switchNode.setThreadName(Thread.currentThread().getName());
        switchNode.setServiceName("THREAD_SWITCH");
        switchNode.setMethodName(Thread.currentThread().getName());
        switchNode.setDepth(parentDepth + 1);
        switchNode.setStartTime(System.currentTimeMillis());
        switchNode.setSuccess(true);
        // 关键：调用TraceContext的统一存储接口
//        TraceContext traceContext = SpringContextHolder.getBeanSafe(TraceContext.class);
//        if (traceContext != null) {
//            traceContext.addNode(switchNode);
//        }
        return switchNode;
    }

    /**
     * 生成线程切换节点的spanId（确保与父节点关联）
     */
    // 新增轻量SpanId生成方法（比UUID快3倍+内存占用减少60%）
    private String generateLightweightSpanId(String parentSpanId) {
        String randomPart = Integer.toHexString(new Random().nextInt(0xFFFFFF)); // 6位随机十六进制数
        if(StringUtils.isBlank(parentSpanId))
        {
            return "ts-" + randomPart;
        }
        else
        {
            return parentSpanId + "-ts-" + randomPart;
        }
    }

    /**
     * 批量包装任务
     */
    private <T> Collection<? extends Callable<T>> wrapTasks(Collection<? extends Callable<T>> tasks) {
        return tasks.stream().map(this::wrapTask).collect(Collectors.toList());
    }

    // -------------------------- 代理原始线程池方法（保持不变） --------------------------
    @Override
    public void shutdown() {
        delegate.shutdown();
    }

    @Override
    public List<Runnable> shutdownNow() {
        return delegate.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return delegate.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return delegate.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return delegate.awaitTermination(timeout, unit);
    }

    @Override
    public int getActiveCount() {
        return delegate.getActiveCount();
    }

    @Override
    public long getTaskCount() {
        return delegate.getTaskCount();
    }
}
